using Avalonia;
using Avalonia.Controls;
using Avalonia.Media.Imaging;
using Avalonia.Platform;
using CommunityToolkit.Mvvm.ComponentModel;
using CommunityToolkit.Mvvm.Input;
using System;
using System.Collections.ObjectModel;
using System.Collections.Generic;
using System.IO;
using System.Runtime.InteropServices;
using System.Threading.Tasks;
using Avalonia.Controls.ApplicationLifetimes;
using System.Linq;
using Avalonia.Platform.Storage;

namespace MathPlus.Features.ScreenShot
{
    public partial class ScreenShotViewModel : ObservableObject
    {
        [ObservableProperty]
        private Bitmap? currentScreenshot;

        [ObservableProperty]
        private ObservableCollection<Screen> availableScreens;

        [ObservableProperty]
        private Screen? selectedScreen;

        [ObservableProperty]
        private bool hasScreenshot;

        public ScreenShotViewModel()
        {
            availableScreens = new ObservableCollection<Screen>();
            UpdateAvailableScreens();
        }

        private void UpdateAvailableScreens()
        {
            AvailableScreens.Clear();
            if (Application.Current?.ApplicationLifetime is IClassicDesktopStyleApplicationLifetime desktop)
            {
                var screens = desktop.MainWindow?.Screens.All ?? Enumerable.Empty<Screen>();
                foreach (var screen in screens)
                {
                    AvailableScreens.Add(screen);
                }
                SelectedScreen = AvailableScreens.Count > 0 ? AvailableScreens[0] : null;
            }
        }

        [RelayCommand]
        private async Task Capture()
        {
            if (SelectedScreen == null) return;

            if (Application.Current?.ApplicationLifetime is IClassicDesktopStyleApplicationLifetime desktop)
            {
                var window = desktop.MainWindow;
                if (window != null)
                {
                    window.WindowState = WindowState.Minimized;
                    await Task.Delay(500); // Wait for window to minimize
                }

                try
                {
                    CurrentScreenshot = await CaptureScreen(SelectedScreen);
                    HasScreenshot = CurrentScreenshot != null;
                }
                finally
                {
                    if (window != null)
                    {
                        window.WindowState = WindowState.Normal;
                    }
                }
            }
        }

        [RelayCommand]
        private async Task Save()
        {
            if (CurrentScreenshot == null) return;

            if (Application.Current?.ApplicationLifetime is IClassicDesktopStyleApplicationLifetime desktop)
            {
                var window = desktop.MainWindow;
                if (window == null) return;

                var options = new FilePickerSaveOptions
                {
                    Title = "Save Screenshot",
                    DefaultExtension = "png",
                    FileTypeChoices = new[] 
                    {
                        new FilePickerFileType("PNG Image")
                        {
                            Patterns = new[] { "*.png" },
                            MimeTypes = new[] { "image/png" }
                        }
                    }
                };

                var file = await window.StorageProvider.SaveFilePickerAsync(options);
                if (file != null)
                {
                    await using var stream = await file.OpenWriteAsync();
                    CurrentScreenshot.Save(stream);
                }
            }
        }

        private async Task<Bitmap?> CaptureScreen(Screen screen)
        {
            if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
            {
                return await CaptureWindowsScreen(screen);
            }
            else if (RuntimeInformation.IsOSPlatform(OSPlatform.Linux))
            {
                return await CaptureLinuxScreen(screen);
            }
            else if (RuntimeInformation.IsOSPlatform(OSPlatform.OSX))
            {
                return await CaptureMacScreen(screen);
            }
            
            return null;
        }

        private async Task<Bitmap> CaptureWindowsScreen(Screen screen)
        {
            var screenWidth = (int)screen.Bounds.Width;
            var screenHeight = (int)screen.Bounds.Height;
            var x = (int)screen.Bounds.X;
            var y = (int)screen.Bounds.Y;

            using var bitmap = new WriteableBitmap(new PixelSize(screenWidth, screenHeight), 
                new Vector(96, 96), PixelFormat.Bgra8888, AlphaFormat.Premul);
            
            using (var context = bitmap.Lock())
            {
                var hdc = Win32.GetDC(IntPtr.Zero);
                var memDc = Win32.CreateCompatibleDC(hdc);
                var hBitmap = Win32.CreateCompatibleBitmap(hdc, screenWidth, screenHeight);
                var oldBitmap = Win32.SelectObject(memDc, hBitmap);

                Win32.BitBlt(memDc, 0, 0, screenWidth, screenHeight, hdc, x, y, Win32.SRCCOPY);

                Win32.SelectObject(memDc, oldBitmap);
                Win32.DeleteDC(memDc);
                Win32.ReleaseDC(IntPtr.Zero, hdc);

                Win32.GetBitmapBits(hBitmap, context.Size.Width * context.Size.Height * 4, context.Address);

                Win32.DeleteObject(hBitmap);
            }

            return bitmap;
        }

        private async Task<Bitmap> CaptureLinuxScreen(Screen screen)
        {
            if (!RuntimeInformation.IsOSPlatform(OSPlatform.Linux))
                throw new PlatformNotSupportedException("This method is only supported on Linux.");

            var screenWidth = (int)screen.Bounds.Width;
            var screenHeight = (int)screen.Bounds.Height;
            var x = (int)screen.Bounds.X;
            var y = (int)screen.Bounds.Y;

            var display = X11.XOpenDisplay(IntPtr.Zero);
            if (display == IntPtr.Zero)
                throw new Exception("Failed to open X display");

            try
            {
                var root = X11.XDefaultRootWindow(display);
                var image = X11.XGetImage(display, root, x, y, screenWidth, screenHeight, 0xFFFFFFFF, 2); // ZPixmap = 2
                if (image == IntPtr.Zero)
                    throw new Exception("Failed to capture screen");

                try
                {
                    using var bitmap = new WriteableBitmap(new PixelSize(screenWidth, screenHeight),
                        new Vector(96, 96), PixelFormat.Bgra8888, AlphaFormat.Premul);

                    using (var locked = bitmap.Lock())
                    {
                        unsafe
                        {
                            var dest = (byte*)locked.Address;

                            for (var row = 0; row < screenHeight; row++)
                            {
                                for (var col = 0; col < screenWidth; col++)
                                {
                                    var pixel = X11.XGetPixel(image, col, row);
                                    var destPixel = dest + (row * locked.RowBytes) + (col * 4);

                                    // Extract BGRA components (assuming 32-bit color)
                                    destPixel[2] = (byte)((pixel >> 16) & 0xFF); // R
                                    destPixel[1] = (byte)((pixel >> 8) & 0xFF);  // G
                                    destPixel[0] = (byte)(pixel & 0xFF);         // B
                                    destPixel[3] = (byte)0xFF;                   // A (full opacity)
                                }
                            }
                        }
                    }

                    return bitmap;
                }
                finally
                {
                    X11.XDestroyImage(image);
                }
            }
            finally
            {
                X11.XCloseDisplay(display);
            }
        }

        private async Task<Bitmap> CaptureMacScreen(Screen screen)
        {
            throw new NotImplementedException("MacOS screen capture is not implemented yet.");
        }
    }

    internal static class Win32
    {
        public const int SRCCOPY = 0x00CC0020;

        [DllImport("gdi32.dll")]
        public static extern bool BitBlt(IntPtr hdc, int x, int y, int cx, int cy, IntPtr hdcSrc, int x1, int y1, int rop);

        [DllImport("gdi32.dll")]
        public static extern IntPtr CreateCompatibleDC(IntPtr hdc);

        [DllImport("gdi32.dll")]
        public static extern IntPtr CreateCompatibleBitmap(IntPtr hdc, int nWidth, int nHeight);

        [DllImport("gdi32.dll")]
        public static extern IntPtr SelectObject(IntPtr hdc, IntPtr hgdiobj);

        [DllImport("gdi32.dll")]
        public static extern bool DeleteDC(IntPtr hdc);

        [DllImport("gdi32.dll")]
        public static extern bool DeleteObject(IntPtr hObject);

        [DllImport("user32.dll")]
        public static extern IntPtr GetDC(IntPtr hwnd);

        [DllImport("user32.dll")]
        public static extern int ReleaseDC(IntPtr hwnd, IntPtr hdc);

        [DllImport("gdi32.dll")]
        public static extern int GetBitmapBits(IntPtr hbmp, int cbBuffer, IntPtr lpvBits);
    }

    internal static class X11
    {
        [DllImport("libX11", EntryPoint = "XOpenDisplay")]
        public static extern IntPtr XOpenDisplay(IntPtr display);

        [DllImport("libX11", EntryPoint = "XCloseDisplay")]
        public static extern int XCloseDisplay(IntPtr display);

        [DllImport("libX11", EntryPoint = "XDefaultRootWindow")]
        public static extern IntPtr XDefaultRootWindow(IntPtr display);

        [DllImport("libX11", EntryPoint = "XGetImage")]
        public static extern IntPtr XGetImage(IntPtr display, IntPtr drawable, int x, int y,
            int width, int height, long plane_mask, int format);

        [DllImport("libX11", EntryPoint = "XDestroyImage")]
        public static extern int XDestroyImage(IntPtr ximage);

        [DllImport("libX11")]
        public static extern IntPtr XGetImageData(IntPtr image);

        [DllImport("libX11")]
        public static extern ulong XGetPixel(IntPtr ximage, int x, int y);

        [StructLayout(LayoutKind.Sequential)]
        public struct XImage
        {
            public int width, height;
            public int xoffset;
            public int format;
            public IntPtr data;
            public int byte_order;
            public int bitmap_unit;
            public int bitmap_bit_order;
            public int bitmap_pad;
            public int depth;
            public int bytes_per_line;
            public int bits_per_pixel;
            public ulong red_mask;
            public ulong green_mask;
            public ulong blue_mask;
        }
    }
}
